import os  # 导入操作系统相关模块
import os.path  # 导入路径操作相关模块
import numpy as np  # 导入NumPy库，用于数值计算
import glob  # 导入glob模块，用于文件匹配
import SimpleITK as sitk  # 导入SimpleITK库，用于医学图像处理


def Resampling(image_path, save_folder, lable=False):
    # 如果保存文件夹不存在，则创建
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)
    # 获取匹配图像路径列表
    filelist = sorted(glob.glob(image_path))
    print(filelist)
    # 遍历图像列表
    for file in filelist:
        print(file)
        # 读取医学图像
        img_itk = sitk.ReadImage(file)
        # 获取原始图像尺寸和间距
        original_size = img_itk.GetSize()
        original_spacing = img_itk.GetSpacing()
        print("original_size:", original_size)
        print("original_spacing:", original_spacing)

        # 计算新的图像尺寸和间距
        new_size = [512, 416, original_size[2]]
        new_spacing = [original_size[0] * (original_spacing[0] / new_size[0]),
                       original_size[1] * (original_spacing[1] / new_size[1]),
                       original_size[2] * (original_spacing[2] / new_size[2])]
        print("new_size:", new_size)
        print("new_spacing:", new_spacing)

        # 初始化重采样滤波器
        resampleSliceFilter = sitk.ResampleImageFilter()
        # 如果不是标签图像
        if lable == False:
            resampleSliceFilter.SetOutputSpacing(new_spacing)
            resampleSliceFilter.SetSize(new_size)
            resampleSliceFilter.SetOutputDirection(img_itk.GetDirection())
            resampleSliceFilter.SetOutputOrigin(img_itk.GetOrigin())
            resampleSliceFilter.SetTransform(sitk.Transform())
            resampleSliceFilter.SetDefaultPixelValue(img_itk.GetPixelIDValue())
            resampleSliceFilter.SetInterpolator(sitk.sitkBSpline)
            # 执行重采样
            resampleImage = resampleSliceFilter.Execute(img_itk)
            resampleImageArray = sitk.GetArrayFromImage(resampleImage)
            resampleImageArray[resampleImageArray < 0] = 0  # 将图像中小于0的元素置为0
        else:  # 对于标签，应使用线性插值以确保原始和重采样后的标签相同
            resampleImage = resampleSliceFilter.Execute(img_itk, new_size, sitk.Transform(), sitk.sitkLinear,
                                                        img_itk.GetOrigin(), new_spacing, img_itk.GetDirection(), 0,
                                                        img_itk.GetPixelIDValue())
            resampleImageArray = sitk.GetArrayFromImage(resampleImage)

        # 从NumPy数组创建SimpleITK图像
        img_itk_new = sitk.GetImageFromArray(resampleImageArray.astype(np.float32))
        img_itk_new.SetSpacing(new_spacing)
        img_itk_new.SetDirection(img_itk.GetDirection())
        img_itk_new.SetOrigin(img_itk.GetOrigin())

        # 生成新的文件名并保存图像
        file_name = os.path.basename(file)
        file_name = file_name.replace('copy', 'resample')
        print("file_name:", file_name)
        save_path = os.path.join(save_folder, file_name)

        # 保存图像
        sitk.WriteImage(img_itk_new, save_path)


if __name__ == '__main__':
    # 调用Resampling函数对医学图像进行重采样，并保存结果
    Resampling(image_path=r'input.nii',
               save_folder=r'./MRI')

"""
# 其他暂时不用的函数
def RestoreToOriginalSize(image_path, save_folder, label=False):
    # 如果保存文件夹不存在，则创建
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)

    # 获取匹配图像路径列表
    filelist = sorted(glob.glob(image_path))

    # 遍历图像列表
    for file in filelist:
        print(file)

        # 读取医学图像
        img_itk = sitk.ReadImage(file)

        # 获取原始图像尺寸和间距
        original_size = img_itk.GetSize()
        original_spacing = img_itk.GetSpacing()

        # 计算新的图像尺寸和间距
        restore_size = (256, 208, original_size[2])
        new_spacing = [original_size[0] * (original_spacing[0] / restore_size[0]),
                       original_size[1] * (original_spacing[1] / restore_size[1]),
                       original_size[2] * (original_spacing[2] / restore_size[2])]

        # 初始化重采样滤波器
        resampleSliceFilter = sitk.ResampleImageFilter()

        # 设置输出尺寸和间距
        resampleSliceFilter.SetOutputSpacing(new_spacing)
        resampleSliceFilter.SetSize(restore_size)
        resampleSliceFilter.SetOutputDirection(img_itk.GetDirection())
        resampleSliceFilter.SetOutputOrigin(img_itk.GetOrigin())
        resampleSliceFilter.SetTransform(sitk.Transform())
        resampleSliceFilter.SetDefaultPixelValue(255)  # 设置为白色填充
        resampleSliceFilter.SetInterpolator(sitk.sitkNearestNeighbor)  # 使用最近邻插值

        # 执行重采样
        Resampleimage = resampleSliceFilter.Execute(img_itk)

        # 生成新的文件名并保存图像
        file_name = os.path.basename(file)
        file_name = file_name.replace('resample', 'restore')
        print("file_name:", file_name)
        sitk.WriteImage(Resampleimage, os.path.join(save_folder, file_name))


def adjust_hu(image_path, save_folder):
    # 如果保存文件夹不存在，则创建
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)
    # 获取匹配图像路径列表
    filelist = sorted(glob.glob(image_path))
    print(filelist)
    # 遍历图像列表
    for file in filelist:
        print(file)
        # 读取医学图像
        img_itk = sitk.ReadImage(file)
        img_array = sitk.GetArrayFromImage(img_itk)
        print("min, max:", np.min(img_array), np.max(img_array))
        # 将像素值限制在0到500之间
        img_array[img_array > 500] = 500
        img_array[img_array < 0] = 0
        # 从NumPy数组创建SimpleITK图像
        img_itk_new = sitk.GetImageFromArray(img_array.astype(np.float32))
        img_itk_new.SetSpacing(img_itk.GetSpacing())
        img_itk_new.SetDirection(img_itk.GetDirection())
        img_itk_new.SetOrigin(img_itk.GetOrigin())
        # 生成新的文件名并保存图像
        file_name = os.path.basename(file)
        file_name = file_name.replace('resample', 'hu500')
        print("file_name:", file_name)
        sitk.WriteImage(img_itk_new, save_folder + file_name)


def adjust_coordinate(image_path, label_path):
    # 获取匹配图像路径列表
    filelist = sorted(glob.glob(image_path))
    print(filelist)
    # 遍历图像列表
    for file in filelist:
        print(file)
        # 读取医学图像
        img_itk = sitk.ReadImage(file)
        # 获取对应的标签图像文件名
        label_file_name = os.path.basename(file).replace('hu', 'liver')
        # 读取标签图像
        label_itk = sitk.ReadImage(label_path + label_file_name)
        # 将标签图像的方向和原始图像相同
        label_itk.SetDirection(img_itk.GetDirection())
        label_itk.SetOrigin(img_itk.GetOrigin())
        # 保存调整后的标签图像
        sitk.WriteImage(label_itk, label_path + 'ad' + label_file_name)
"""
